#!/usr/bin/env python3

import os
import sys
script_root = os.path.dirname(os.path.realpath(__file__))
sys.path.append(f'{script_root}/../../../tools')
from gmcore_plot import *

zlev = np.linspace(0, 20, 51) * units.km
fig = plt.figure(figsize=(8, 8))

f = xr.open_dataset(f'{script_root}/sc.360x181x40.h0.nc')

z = f.gz.isel(time=0).sel(lon=0) * units.m**2 / units.s**2 / g
z.attrs['units'] = 'm'

# ------------------------------------------------------------------------------
levels = np.linspace(-17.5, 17.5, 15)
colors = ListedColormap([
  [0.00, 0.00, 1.00],
  [0.00, 0.83, 1.00],
  [0.09, 1.00, 1.00],
  [0.60, 1.00, 1.00],
  [0.40, 1.00, 0.60],
  [0.20, 1.00, 0.00],
  [0.80, 1.00, 0.00],
  [0.87, 0.96, 0.00],
  [1.00, 1.00, 0.00],
  [1.00, 0.80, 0.00],
  [1.00, 0.60, 0.00],
  [1.00, 0.40, 0.00],
  [1.00, 0.20, 0.00],
  [1.00, 0.00, 0.00],
])
ax = fig.add_subplot(3, 2, 1)
u = vinterp_z(z, f.u.isel(time=0).sel(lon=0), zlev)
plot_contour_zonal_height(ax, u, cmap=colors, levels=levels, ticks=levels[1::2], with_contour=True, with_grid=False, cbar_orient='horizontal')

# ------------------------------------------------------------------------------
levels = np.linspace(0, 16, 9)
colors = ListedColormap([
  [0.00, 0.00, 1.00],
  [0.09, 1.00, 1.00],
  [0.40, 1.00, 0.60],
  [0.49, 1.00, 0.00],
  [1.00, 1.00, 0.00],
  [1.00, 0.80, 0.00],
  [1.00, 0.40, 0.00],
  [1.00, 0.00, 0.00],
])
ax = fig.add_subplot(3, 2, 2)
q = vinterp_z(z, f.q_qv.isel(time=0).sel(lon=0), zlev)
q = q / (1 + q) * 1000
q.attrs['units'] = 'g kg-1'
plot_contour_zonal_height(ax, q, cmap=colors, levels=levels, with_contour=True, with_grid=False, cbar_orient='horizontal')

# ------------------------------------------------------------------------------
levels = [290, 300, 310, 320, 330, 340, 350, 360, 380, 400, 420, 440, 460, 480, 500, 520]
colors = ListedColormap([
  [0.00, 0.00, 1.00],
  [0.00, 0.83, 1.00],
  [0.09, 1.00, 1.00],
  [0.60, 1.00, 1.00],
  [0.40, 1.00, 0.60],
  [0.20, 1.00, 0.00],
  [0.49, 1.00, 0.00],
  [0.87, 0.96, 0.00],
  [1.00, 1.00, 0.00],
  [1.00, 0.91, 0.00],
  [1.00, 0.80, 0.00],
  [1.00, 0.50, 0.00],
  [1.00, 0.40, 0.00],
  [1.00, 0.16, 0.00],
  [1.00, 0.00, 0.00],
])
norm = BoundaryNorm(levels, ncolors=colors.N, clip=True)
ax = fig.add_subplot(3, 2, 3)
ptv = calc.virtual_potential_temperature(f.p.isel(time=0).sel(lon=0), f.t.isel(time=0).sel(lon=0), f.q_qv.isel(time=0).sel(lon=0))
ptv = vinterp_z(z, ptv, zlev)
plot_contour_zonal_height(ax, ptv, cmap=colors, norm=norm, with_contour=True, with_grid=False, cbar_orient='horizontal')

# ------------------------------------------------------------------------------
levels = np.linspace(-2.2, 2.2, 23)
colors = ListedColormap([
  [0.00, 0.00, 0.50],
  [0.00, 0.00, 0.60],
  [0.00, 0.00, 0.70],
  [0.00, 0.00, 0.78],
  [0.00, 0.00, 0.88],
  [0.00, 0.00, 0.98],
  [0.16, 0.16, 1.00],
  [0.32, 0.32, 1.00],
  [0.52, 0.52, 1.00],
  [0.72, 0.72, 1.00],
  [0.92, 0.92, 1.00],
  [1.00, 0.92, 0.92],
  [1.00, 0.72, 0.72],
  [1.00, 0.52, 0.52],
  [1.00, 0.32, 0.32],
  [1.00, 0.16, 0.16],
  [0.98, 0.00, 0.00],
  [0.88, 0.00, 0.00],
  [0.78, 0.00, 0.00],
  [0.70, 0.00, 0.00],
  [0.60, 0.00, 0.00],
  [0.50, 0.00, 0.00],
])
ax = fig.add_subplot(3, 2, 4)
dptv = ptv[:,90:] - ptv[:,90]
plot_contour_zonal_height(ax, dptv, cmap=colors, levels=levels, ticks=levels[2:-2:2], with_contour=True, with_grid=False, cbar_orient='horizontal')

# ------------------------------------------------------------------------------
levels = [0, 10, 25, 50, 100, 200, 300, 400, 500, 600, 700, 800, 900, 1000, 1100]
colors = ListedColormap([
  [0.00, 0.00, 1.00],
  [0.00, 0.83, 1.00],
  [0.09, 1.00, 1.00],
  [0.60, 1.00, 1.00],
  [0.40, 1.00, 0.60],
  [0.20, 1.00, 0.00],
  [0.80, 1.00, 0.00],
  [0.87, 0.96, 0.00],
  [1.00, 1.00, 0.00],
  [1.00, 0.80, 0.00],
  [1.00, 0.60, 0.00],
  [1.00, 0.40, 0.00],
  [1.00, 0.20, 0.00],
  [1.00, 0.00, 0.00],
])
norm = BoundaryNorm(levels, ncolors=colors.N, clip=True)
ax = fig.add_subplot(3, 2, 5)
p = vinterp_z(z, f.p.isel(time=0).sel(lon=0), zlev)
plot_contour_zonal_height(ax, p / 100, cmap=colors, norm=norm, with_contour=True, with_grid=False, cbar_orient='horizontal')

# ------------------------------------------------------------------------------
levels = np.linspace(-1.1, 1.1, 23)
colors = ListedColormap([
  [0.00, 0.00, 0.50],
  [0.00, 0.00, 0.60],
  [0.00, 0.00, 0.70],
  [0.00, 0.00, 0.78],
  [0.00, 0.00, 0.88],
  [0.00, 0.00, 0.98],
  [0.16, 0.16, 1.00],
  [0.32, 0.32, 1.00],
  [0.52, 0.52, 1.00],
  [0.72, 0.72, 1.00],
  [0.92, 0.92, 1.00],
  [1.00, 0.92, 0.92],
  [1.00, 0.72, 0.72],
  [1.00, 0.52, 0.52],
  [1.00, 0.32, 0.32],
  [1.00, 0.16, 0.16],
  [0.98, 0.00, 0.00],
  [0.88, 0.00, 0.00],
  [0.78, 0.00, 0.00],
  [0.70, 0.00, 0.00],
  [0.60, 0.00, 0.00],
  [0.50, 0.00, 0.00],
])
ax = fig.add_subplot(3, 2, 6)
dp = p[:,90:] - p[:,90]
plot_contour_zonal_height(ax, dp / 100, cmap=colors, levels=levels, ticks=levels[2:-2:2], with_contour=True, with_grid=False, cbar_orient='horizontal')

plt.tight_layout()
plt.show()
